
import argparse
import os
import json
import cv2
import numpy as np
from annotation_stat import open_spire_annotations


def main():
    parser = argparse.ArgumentParser(description="Statistic on spire annotations")
    parser.add_argument(
        "--spire-dir",
        default="C:/dataset/spire_dataset/BB-DC-drone-bird-trainval",
        help="path to image dir",
    )
    args = parser.parse_args()

    img_dir = os.path.join(args.spire_dir, 'scaled_images')
    ann_dir = os.path.join(args.spire_dir, 'annotations')

    input_names = os.listdir(img_dir)
    for img_name in input_names:

        if img_name.endswith('.jpg') or img_name.endswith('.jpeg') or img_name.endswith('.png'):
            img_fn = os.path.join(img_dir, img_name)
            img = cv2.imread(img_fn)

            img_name_base = os.path.splitext(img_name)[0]

            json_f = open(os.path.join(ann_dir, img_name + '.json'), 'r')
            json_str = json_f.read()
            json_dict = json.loads(json_str)
            json_f.close()

            new_img_name = img_name_base + '.jpg'
            json_dict['file_name'] = new_img_name

            cv2.imwrite(os.path.join(img_dir, new_img_name), img)
            os.remove(img_fn)

            fp = open(os.path.join(ann_dir, json_dict['file_name']+'.json'), 'w')
            json.dump(json_dict, fp)
            if img_name != new_img_name:
                os.remove(os.path.join(ann_dir, img_name + '.json'))


if __name__ == '__main__':
    main()
